import cv2
import numpy as np
import yaml

def load_calibration_parameters(calibration_file='camera_calibration1280.yaml'):
    """
    加载相机标定参数
    """
    try:
        with open(calibration_file, 'r') as f:
            calibration_data = yaml.safe_load(f)
        
        camera_matrix = np.array(calibration_data['camera_matrix'])
        dist_coeffs = np.array(calibration_data['distortion_coefficients'])
        
        print("成功加载相机标定参数")
        print(f"标定日期: {calibration_data['calibration_date']}")
        print(f"重投影误差: {calibration_data['reprojection_error']:.6f} 像素")
        print(f"使用图片数量: {calibration_data['image_count']}")
        
        return camera_matrix, dist_coeffs, calibration_data
    
    except FileNotFoundError:
        print(f"错误：找不到标定文件 '{calibration_file}'")
        print("请先运行标定程序生成 camera_calibration.yaml 文件")
        exit()
    except Exception as e:
        print(f"加载标定文件时出错: {e}")
        exit()

def setup_undistort_maps(camera_matrix, dist_coeffs, resolution):
    """
    设置矫正映射，提高实时处理性能
    """
    # 获取优化的新相机矩阵和ROI区域
    new_camera_matrix, roi = cv2.getOptimalNewCameraMatrix(
        camera_matrix, dist_coeffs, resolution, 1, resolution
    )
    
    # 生成矫正映射
    map_x, map_y = cv2.initUndistortRectifyMap(
        camera_matrix, dist_coeffs, None, new_camera_matrix, resolution, cv2.CV_32FC1
    )
    
    return map_x, map_y, roi

def main():
    # 加载标定参数
    camera_matrix, dist_coeffs, calibration_data = load_calibration_parameters()
    
    # 打开摄像头
    cap = cv2.VideoCapture(0)
    
    if not cap.isOpened():
        print("错误：无法打开摄像头")
        exit()
    
    # 设置摄像头分辨率（与标定时相同）
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    
    # 获取图像分辨率
    ret, test_frame = cap.read()
    if not ret:
        print("错误：无法从摄像头读取图像")
        cap.release()
        exit()
    
    resolution = (test_frame.shape[1], test_frame.shape[0])  # (width, height)
    
    # 设置矫正映射
    map_x, map_y, roi = setup_undistort_maps(camera_matrix, dist_coeffs, resolution)
    x, y, w, h = roi
    
    print("\n=== 实时矫正对比程序 ===")
    print("按键说明:")
    print("  'q' - 退出程序")
    print("  's' - 保存当前对比图像")
    print("  'f' - 切换全屏/窗口模式")
    print("  'r' - 重置ROI裁剪")
    print("\n红色边框: 原始图像 (有畸变)")
    print("绿色边框: 矫正后图像 (无畸变)")
    
    fullscreen = False
    show_roi_crop = True
    
    while True:
        # 读取帧
        ret, frame = cap.read()
        if not ret:
            print("错误：无法读取帧")
            break
        
        # 使用映射进行快速矫正
        undistorted_frame = cv2.remap(frame, map_x, map_y, cv2.INTER_LINEAR)
        
        # 裁剪ROI区域（去除黑边）
        if show_roi_crop and w > 0 and h > 0:
            cropped_frame = undistorted_frame[y:y+h, x:x+w]
        else:
            cropped_frame = undistorted_frame
        
        # 调整图像大小以便并排显示
        orig_display = cv2.resize(frame, (320, 240))
        undistorted_display = cv2.resize(cropped_frame, (320, 240))
        
        # 添加边框和文字标注
        orig_with_border = cv2.copyMakeBorder(orig_display, 40, 10, 10, 10, cv2.BORDER_CONSTANT, value=(0, 0, 255))
        undistorted_with_border = cv2.copyMakeBorder(undistorted_display, 40, 10, 10, 10, cv2.BORDER_CONSTANT, value=(0, 255, 0))
        
        # 添加文字说明
        cv2.putText(orig_with_border, "Original (With Distortion)", (10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
        cv2.putText(undistorted_with_border, "Undistorted (Corrected)", (10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
        
        # 添加分辨率信息
        cv2.putText(orig_with_border, f"{resolution[0]}x{resolution[1]}", (10, 270), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
        cv2.putText(undistorted_with_border, f"{w}x{h}" if show_roi_crop else f"{resolution[0]}x{resolution[1]}", (10, 270), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
        
        # 并排显示
        comparison = np.hstack((orig_with_border, undistorted_with_border))
        
        # 添加中间分隔线
        cv2.line(comparison, (330, 0), (330, 290), (255, 255, 255), 2)
        
        # 显示图像
        window_name = "Camera Distortion Correction Comparison"
        if fullscreen:
            cv2.namedWindow(window_name, cv2.WND_PROP_FULLSCREEN)
            cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
        else:
            cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
            cv2.resizeWindow(window_name, 700, 350)
        
        cv2.imshow(window_name, comparison)
        
        # 键盘控制
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord('s'):
            # 保存对比图像
            timestamp = cv2.getTickCount()
            filename = f"comparison_{timestamp}.jpg"
            cv2.imwrite(filename, comparison)
            print(f"已保存对比图像: {filename}")
        elif key == ord('f'):
            # 切换全屏模式
            fullscreen = not fullscreen
            print(f"全屏模式: {'开启' if fullscreen else '关闭'}")
        elif key == ord('r'):
            # 切换ROI裁剪
            show_roi_crop = not show_roi_crop
            print(f"ROI裁剪: {'开启' if show_roi_crop else '关闭'}")
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()
    print("程序已退出")

if __name__ == "__main__":
    main()